package com.earthea.common.algorithm.dictionarytree;

import com.google.common.base.Preconditions;
import org.apache.commons.lang3.tuple.Pair;

import java.util.List;
import java.util.Optional;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;

/**
 * 字典树算法
 *
 * @param <T> 节点路径类型
 * @param <D> 节点数据类型
 */
public class DictionaryTreeAlgorithmImpl<T, X, D> implements DictionaryTreeAlgorithm<T, D> {
    private Function<T, List<X>> nodeItemProvider;

    private Optional<Predicate<X>> variablePredictor;

    private Supplier<X> variableSupplier = null;

    private TreeNode<X, D> rootNode = new TreeNode<>();

    public DictionaryTreeAlgorithmImpl(Function<T, List<X>> nodeItemProvider, Predicate<X> variablePredictor, Supplier<X> variableSupplier) {
        this.nodeItemProvider = Preconditions.checkNotNull(nodeItemProvider, "nodeItemProvider should be not null");
        this.variablePredictor = Optional.ofNullable(variablePredictor);
        if (variableSupplier != null) {
            this.variableSupplier = variableSupplier;
        }
    }

    public DictionaryTreeAlgorithmImpl(Function<T, List<X>> nodeItemProvider) {
        this(nodeItemProvider, null, null);
    }

    @Override
    public void addNode(T nodePath, boolean isPrefix, D nodeData) {
        final List<X> nodeItems = nodeItemProvider.apply(nodePath);
        TreeNode<X, D> curNode = rootNode;
        for (X nodeItem : nodeItems) {
            Pair<Boolean, X> isVariable = isVariable(nodeItem);
            X normalizeItem = normalizeItem(nodeItem, isVariable);
            curNode = curNode.getChildNodes().computeIfAbsent(normalizeItem, k -> new TreeNode<>(normalizeItem, isVariable.getLeft()));
        }
        curNode.setMatch(true);
        curNode.setData(nodeData);
        curNode.setPrefixMatch(curNode.isPrefixMatch() | isPrefix);
    }

    private Pair<Boolean, X> isVariable(X nodeItem) {
        boolean isVariable = variablePredictor.map(predicate -> predicate.test(nodeItem)).orElse(false);
        return isVariable ? Pair.of(true, variableSupplier.get()) : Pair.of(false, nodeItem);
    }

    private X normalizeItem(X nodeItem, Pair<Boolean, X> isVariable) {
        return isVariable.getLeft() ? isVariable.getRight() : nodeItem;
    }

    @Override
    public boolean isMatch(T nodePath) {
        return getMatchData(nodeItemProvider.apply(nodePath), 0, rootNode) != null;
    }

    @Override
    public D getMatchData(T nodePath) {
        return Optional.ofNullable(getMatchData(nodeItemProvider.apply(nodePath), 0, rootNode)).map(TreeNode::getData).orElse(null);
    }

    private TreeNode<X, D> getMatchData(List<X> items, int index, TreeNode<X, D> parentNode) {
        if (parentNode == null || index >= items.size()) {
            return null;
        }
        X item = items.get(index);
        TreeNode<X, D> variableNode = variableSupplier != null ? parentNode.getChildNodes().get(variableSupplier.get()) : null;
        TreeNode<X, D> curNode = Optional.ofNullable(parentNode.getChildNodes().get(item)).orElse(variableNode);
        if (curNode == null) {
            return null;
        }
        if (curNode.isPrefixMatch() || (index == items.size() - 1 && curNode.isMatch())) {
            return curNode;
        }
        TreeNode<X, D> matchData = getMatchData(items, index + 1, curNode);
        if (matchData != null) {
            return matchData;
        }
        return curNode != variableNode ? getMatchData(items, index + 1, variableNode) : null;
    }
}
